# This is a diameter message in bytes-like format.
# In python3, from socket object, you received message in bytes type like below.
import dpkt as dpkt
from scapy.all import *
from pyDiameter.pyDiaMessage import DiaMessage
from pyDiameter.pyDiaAVPTypes import DiaAVPStr
from pyDiameter.pyDiaAVPPath import DiaAVPPath

from pyDiameter.pyDiaAVPTools import address_to_bytes, bytes_to_address, time_to_bytes, bytes_to_time
from pyDiameter.pyDiaAVPPath import DiaAVPPath
from pyDiameter.pyDiaAVPBasicTypes import DiaAVPStr


def get_data(path):
    pcap_pkts = rdpcap(path)
    msgBuffs = []
    for packet in pcap_pkts:
        # packet.show()
        # print(packet['SCTPChunkData'])
        decode_data = packet['SCTPChunkData'].data
        # msgBuffs.append(decode_data)
        msgBuffs.append(decode_data.hex())
    return msgBuffs

def get_datas(path):
    fpcap = open(path, 'rb')
    string_data = fpcap.read()
    string_data.hex()
    print()

# msgBuff=b'\x01\x00\x03\xf0\x40\x00\x01\x3c\x01\x00\x00\x23\x63\xb2\x14\x08\x63\xb2\x14\x08\x00\x00\x01\x07\x40\x00\x00\x54\x6d\x6d\x65\x63\x30\x33\x2e\x6d\x6d\x65\x67\x69\x38\x30\x30\x33\x2e\x6d\x6d\x65\x2e\x65\x70\x63\x2e\x6d\x6e\x63\x30\x30\x30\x2e\x6d\x63\x63\x34\x36\x30\x2e\x33\x67\x70\x70\x6e\x65\x74\x77\x6f\x72\x6b\x2e\x6f\x72\x67\x3b\x33\x38\x30\x38\x34\x32\x39\x32\x31\x38\x3b\x31\x33\x36\x34\x34\x39\x3b\x6d\x6d\x65\x00\x00\x01\x15\x40\x00\x00\x0c\x00\x00\x00\x01\x00\x00\x01\x08\x40\x00\x00\x39\x68\x73\x73\x35\x2e\x66\x7a\x2e\x66\x6a\x2e\x6e\x6f\x64\x65\x2e\x69\x6d\x73\x2e\x6d\x6e\x63\x30\x30\x30\x2e\x6d\x63\x63\x34\x36\x30\x2e\x33\x67\x70\x70\x6e\x65\x74\x77\x6f\x72\x6b\x2e\x6f\x72\x67\x00\x00\x00\x00\x00\x01\x28\x40\x00\x00\x29\x69\x6d\x73\x2e\x6d\x6e\x63\x30\x30\x30\x2e\x6d\x63\x63\x34\x36\x30\x2e\x33\x67\x70\x70\x6e\x65\x74\x77\x6f\x72\x6b\x2e\x6f\x72\x67\x00\x00\x00\x00\x00\x05\x78\xc0\x00\x02\x88\x00\x00\x28\xaf\x00\x00\x05\x90\xc0\x00\x00\x10\x00\x00\x28\xaf\x00\x00\x00\x00\x00\x00\x02\xbd\xc0\x00\x00\x13\x00\x00\x28\xaf\x68\x31\x04\x00\x03\x13\xf9\x00\x00\x00\x05\x99\xc0\x00\x00\x11\x00\x00\x28\xaf\x68\x21\x43\x00\x20\x00\x00\x00\x00\x00\x05\x9b\xc0\x00\x00\x2c\x00\x00\x28\xaf\x00\x00\x02\x04\xc0\x00\x00\x10\x00\x00\x28\xaf\x0f\x42\x40\x00\x00\x00\x02\x03\xc0\x00\x00\x10\x00\x00\x28\xaf\x0f\x42\x40\x00\x00\x00\x05\x95\xc0\x00\x02\x18\x00\x00\x28\xaf\x00\x00\x05\x8f\xc0\x00\x00\x10\x00\x00\x28\xaf\x00\x00\x00\x01\x00\x00\x05\x94\xc0\x00\x00\x10\x00\x00\x28\xaf\x00\x00\x00\x00\x00\x00\x05\x96\xc0\x00\x00\xf8\x00\x00\x28\xaf\x00\x00\x05\x8f\xc0\x00\x00\x10\x00\x00\x28\xaf\x00\x00\x00\x01\x00\x00\x05\xb0\xc0\x00\x00\x10\x00\x00\x28\xaf\x00\x00\x00\x00\x00\x00\x01\xed\x40\x00\x00\x0d\x63\x6d\x6e\x65\x74\x00\x00\x00\x00\x00\x05\x97\xc0\x00\x00\x58\x00\x00\x28\xaf\x00\x00\x04\x04\xc0\x00\x00\x10\x00\x00\x28\xaf\x00\x00\x00\x09\x00\x00\x04\x0a\x80\x00\x00\x3c\x00\x00\x28\xaf\x00\x00\x04\x16\x80\x00\x00\x10\x00\x00\x28\xaf\x00\x00\x00\x02\x00\x00\x04\x17\x80\x00\x00\x10\x00\x00\x28\xaf\x00\x00\x00\x00\x00\x00\x04\x18\x80\x00\x00\x10\x00\x00\x28\xaf\x00\x00\x00\x00\x00\x00\x05\x98\xc0\x00\x00\x10\x00\x00\x28\xaf\x00\x00\x00\x01\x00\x00\x01\xe6\x40\x00\x00\x18\x00\x00\x01\x4e\x40\x00\x00\x0e\x00\x01\xc0\xa8\xa6\x95\x00\x00\x00\x00\x05\x9e\xc0\x00\x00\x10\x00\x00\x28\xaf\x00\x00\x00\x00\x00\x00\x05\x9b\xc0\x00\x00\x2c\x00\x00\x28\xaf\x00\x00\x02\x04\xc0\x00\x00\x10\x00\x00\x28\xaf\x0f\x42\x40\x00\x00\x00\x02\x03\xc0\x00\x00\x10\x00\x00\x28\xaf\x0f\x42\x40\x00\x00\x00\x05\x96\xc0\x00\x00\xf4\x00\x00\x28\xaf\x00\x00\x05\x8f\xc0\x00\x00\x10\x00\x00\x28\xaf\x00\x00\x00\x02\x00\x00\x05\xb0\xc0\x00\x00\x10\x00\x00\x28\xaf\x00\x00\x00\x01\x00\x00\x01\xed\x40\x00\x00\x0b\x69\x6d\x73\x00\x00\x00\x05\x97\xc0\x00\x00\x58\x00\x00\x28\xaf\x00\x00\x04\x04\xc0\x00\x00\x10\x00\x00\x28\xaf\x00\x00\x00\x09\x00\x00\x04\x0a\x80\x00\x00\x3c\x00\x00\x28\xaf\x00\x00\x04\x16\x80\x00\x00\x10\x00\x00\x28\xaf\x00\x00\x00\x02\x00\x00\x04\x17\x80\x00\x00\x10\x00\x00\x28\xaf\x00\x00\x00\x00\x00\x00\x04\x18\x80\x00\x00\x10\x00\x00\x28\xaf\x00\x00\x00\x00\x00\x00\x05\x98\xc0\x00\x00\x10\x00\x00\x28\xaf\x00\x00\x00\x01\x00\x00\x01\xe6\x40\x00\x00\x18\x00\x00\x01\x4e\x40\x00\x00\x0e\x00\x01\xc0\xa8\xa6\x95\x00\x00\x00\x00\x05\x9e\xc0\x00\x00\x10\x00\x00\x28\xaf\x00\x00\x00\x00\x00\x00\x05\x9b\xc0\x00\x00\x2c\x00\x00\x28\xaf\x00\x00\x02\x04\xc0\x00\x00\x10\x00\x00\x28\xaf\x0f\x42\x40\x00\x00\x00\x02\x03\xc0\x00\x00\x10\x00\x00\x28\xaf\x0f\x42\x40\x00\x00\x00\x05\x7e\xc0\x00\x00\x10\x00\x00\x28\xaf\x00\x00\x00\x01\x00\x00\x02\x74\x80\x00\x00\x38\x00\x00\x28\xaf\x00\x00\x01\x0a\x40\x00\x00\x0c\x00\x00\x28\xaf\x00\x00\x02\x75\x80\x00\x00\x10\x00\x00\x28\xaf\x00\x00\x00\x01\x00\x00\x02\x76\x80\x00\x00\x10\x00\x00\x28\xaf\x24\x00\x04\x07\x00\x00\x02\x74\x80\x00\x00\x38\x00\x00\x28\xaf\x00\x00\x01\x0a\x40\x00\x00\x0c\x00\x00\x28\xaf\x00\x00\x02\x75\x80\x00\x00\x10\x00\x00\x28\xaf\x00\x00\x00\x02\x00\x00\x02\x76\x80\x00\x00\x10\x00\x00\x28\xaf\x00\x01\x00\x00\x00\x00\x01\x0c\x40\x00\x00\x0c\x00\x00\x07\xd1'
# msgBuff=b'\xc0\x00\x00\x10\x00\x00(\xaf\x00\x00\x00\x01\x00\x00\x01\xf8\xc0\x00\x00H\x00\x00(\xaf+g.3gpp.icsi-ref="urn%3Aurn-7%3A3gpp-service.ims.icsi.mmtel"\x00\x00\x02\x08\xc0\x00\x00\x10\x00\x00(\xaf\x00\x00\x00\x01\x00\x00\x02\x04\xc0\x00\x00\x10\x00\x00(\xaf\x00\x00\x00\x00\x00\x00\x02\x03\xc0\x00\x00\x10\x00\x00(\xaf\x00\x07\xd0\x00\x00\x00\x02\x0f\xc0\x00\x00\x10\x00\x00(\xaf\x00\x00\x00\x01\x00\x00\x01\xf9\xc0\x00\x00Y\x00\x00(\xafsg1.psbc21ber.sh.sh.node.epc.mnc011.mcc460.3gppnetwork.org-1536-308193-798942\x00\x00\x00\x00\x00\x02\x01\xc0\x00\x00\x10\x00\x00(\xaf\x00\x00\x00\x0c\x00\x00\x02t\xc0\x00\x008\x00\x00(\xaf\x00\x00\x01\n@\x00\x00\x0c\x00\x00(\xaf\x00\x00\x02u\x80\x00\x00\x10\x00\x00(\xaf\x00\x00\x00\x01\x00\x00\x02v\x80\x00\x00\x10\x00\x00(\xaf\x00\x00\x003\x00\x00\x00a@\x00\x00\x1a\x00\x80$\x0e\x00k\x00\x01\x022\xc1\x84\xf4\x01F2Y\x93\x00\x00\x00\x00\x02\x18\x80\x00\x00\x10\x00\x00(\xaf\x00\x00\x00\x00\x00\x00\x01\x16@\x00\x00\x0c\x00\x00\x00\x02'


def visitAVP(avp, tab=''):
    print(tab, end='')
    print('name:  ', avp.getAVPName())
    print(tab, end='')
    print('type:  ', avp.getAVPType())
    print(tab, end='')
    print('code:  ', avp.getAVPCode())
    print(tab, end='')
    print('flags: ', avp.getAVPFlags())
    print(tab, end='')
    print('len:   ', len(avp))
    value = avp.getAVPValue()
    if avp.getAVPVSFlag():
        print(tab, end='')
        print('vendor:', avp.getAVPVendor())
    if type(value) is list:
        print(tab, end='')
        print('====>')
        for sub in value:
            visitAVP(sub, tab + '    ')
        print(tab, end='')
        print('<====')
    else:
        print(tab, end='')
        print('value: ', value)
    print(tab, end='')
    print('-------')


def visitMessage(msg):
    print('len:   ', len(msg))
    print('flags: ', msg.getFlags())
    print('code:  ', msg.getCommandCode())
    print('app:   ', msg.getApplicationID())
    print('hbh:   ', msg.getHBHID())
    print('e2e:   ', msg.getE2EID())
    avps = msg.getAVPs()
    for avp in avps:
        visitAVP(avp)


from pyDiameter.pyDiaMessage import *

path = "D:/python/diameter/data/dia_frag.pcap"
msgBuffs = get_data(path)
msgBuff = msgBuffs[0]
avpBuff3 = bytes.fromhex(msgBuff)

msg = DiaMessage()
msg.decode(msgBuff)
visitMessage(msg)

print("=========== Add an address AVP by using DiaAVPPath ===========")
# as we known, avp1 is a group avp, only group avp could add sub-avp.
from pyDiameter.pyDiaAVPBasicTypes import DiaAVPStr
from pyDiameter.pyDiaAVPPath import DiaAVPPath
from pyDiameter.pyDiaAVPTools import address_to_bytes, bytes_to_address, time_to_bytes, bytes_to_time

newAVP = DiaAVPStr()
newAVP.setAVPCode(257)
newAVP.setAVPMandatoryFlag()
newAVP.setAVPValue(address_to_bytes(('ipv6', '2000:1234:ac12:abcd:abcd:abcd:abcd:abc1')))
avpPath = DiaAVPPath()
avpPath.setPath('10415/1400[0]')
msg.addAVPByPath(avpPath, newAVP)
visitMessage(msg)

print("=========== Remove a Feature-List AVP in Second Supported-Feature AVP ===========")
avpPath = DiaAVPPath()
avpPath.setPath('10415/628[1]->10415/630[0]')
msg.removeAVPByPath(avpPath)
visitMessage(msg)

# Try following functions and find the result by yourself.
msg.generateHBHID()
msg.generateE2EID()
msg.setRequestFlag()
msg.clearRequestFlag()
msg.setProxyableFlag()
msg.clearProxyableFlag()
msg.setEBitFlag()
msg.clearEBitFlag()
msg.setTBitFlag()
msg.clearTBitFlag()